from time import sleep
from pymilvus import DataType, connections, utility, Collection, CollectionSchema, FieldSchema, DataType
import os
import numpy

from ..base.module import BaseANN


def metric_mapping(_metric: str):
    _metric_type = {"angular": "COSINE", "euclidean": "L2"}.get(_metric, None)
    if _metric_type is None:
        raise Exception(f"[Milvus] Not support metric type: {_metric}!!!")
    return _metric_type


def get_conn():
    conn = connections.connect("default", host='localhost', port='19530', pool_size=600, retries=3, timeout=10, request_timeout=120)
    return conn


class Milvus(BaseANN):
    def __init__(self, metric, dim, index_param):
        super().__init__()
        self._metric = metric
        self._dim = dim
        self._metric_type = metric_mapping(self._metric)
        self.connects = get_conn()
        self.collection_name = "test_milvus"
        self.collection = None
        self.search_params = None

    @property
    def tablename(self):
        return self.collection_name

    @tablename.setter
    def tablename(self, value):
        self.collection_name = value
        self._update_query()  # 当tablename变化时更新查询语句

    def _update_query(self):
        """根据当前的metric和tablename更新查询语句"""
        # Milvus使用search方法而不是SQL查询，所以这里不需要实现
        pass

    def get_row_count(self) -> int:
        """获取表中的行数"""
        if not self.collection:
            self.collection = Collection(
                self.collection_name,
                consistence_level="STRONG"
            )
        return self.collection.num_entities

    def get_vector_by_id(self, vector_id):
        """根据ID获取向量"""
        try:
            if not self.collection:
                self.collection = Collection(
                    self.collection_name,
                    consistence_level="STRONG"
                )
            
            # 使用ID范围查询，确保能找到精确匹配
            result = self.collection.query(
                expr=f"id == {vector_id}",
                output_fields=["vector"]
            )
            
            if not result or len(result) == 0:
                print(f"[Milvus] 在集合 {self.collection_name} 中找不到ID = {vector_id}的记录")
                return None
            
            # 确保返回的是numpy数组格式
            vector_data = result[0].get("vector")
            if not isinstance(vector_data, numpy.ndarray):
                vector_data = numpy.array(vector_data)
                
            return vector_data
        except Exception as e:
            print(f"[Milvus] 获取向量出错: {e}")
            return None

    def has_same_data(self, dataset) -> bool:
        """检查是否可以重用现有数据"""
        try:
            # 1. 检查集合是否存在
            if not utility.has_collection(self.collection_name):
                print(f"[Milvus] 集合 {self.collection_name} 不存在，无法重用数据")
                return False
                
            # 初始化collection对象
            if not self.collection:
                self.collection = Collection(
                    self.collection_name,
                    consistence_level="STRONG"
                )
            
            # 2. 检查行数是否匹配
            train_data = dataset["train"]
            total_rows = train_data.shape[0]
            db_row_count = self.get_row_count()
            
            # 允许最多2%的数据差异
            max_diff_percentage = 0.02
            diff_rows = abs(db_row_count - total_rows)
            diff_percentage = diff_rows / total_rows if total_rows > 0 else 1.0
            
            if diff_percentage > max_diff_percentage:
                print(f"[Milvus] 数据行数差异过大: 数据库={db_row_count}, 训练集={total_rows}, 差异率={diff_percentage:.2%}")
                return False
            elif diff_rows > 0:
                print(f"[Milvus] 数据行数存在少量差异: 数据库={db_row_count}, 训练集={total_rows}, 差异行数={diff_rows}")
                # 对于小差异，我们可以继续检查
            
            # 3. 检查样本点
            # 选择多个检查点
            check_points = [0]  # 第一个点
            if total_rows > 1:
                check_points.append(total_rows - 1)  # 尝试最后一个点，如果不存在就减少
                # 如果数据库行数比训练集少，调整最后一个检查点
                if db_row_count < total_rows:
                    check_points[-1] = db_row_count - 1
                    
            if total_rows > 10:
                check_points.append(min(total_rows, db_row_count) // 2)  # 中间点
            
            for point_id in check_points:
                try:
                    vector = train_data[point_id]
                    
                    # 获取数据库中的向量
                    db_values = self.get_vector_by_id(point_id)
                    if db_values is None:
                        print(f"[Milvus] 数据库中找不到ID={point_id}的记录")
                        return False
                    
                    # 判断长度是否一致
                    if len(db_values) != len(vector):
                        print(f"[Milvus] ID={point_id}的向量长度不同: 数据库={len(db_values)}, 训练集={len(vector)}")
                        return False
                    
                    # 向量比较 - 使用向量相似度比较
                    import numpy as np
                    
                    # 根据度量类型选择比较方法
                    if self._metric == "angular":
                        # 余弦相似度比较
                        norm1 = np.linalg.norm(vector)
                        norm2 = np.linalg.norm(db_values)
                        if norm1 == 0 or norm2 == 0:
                            similarity = 1.0 if norm1 == norm2 else 0.0
                        else:
                            similarity = np.dot(vector, db_values) / (norm1 * norm2)
                        
                        if similarity < 0.9999:  # 设置一个非常接近1的阈值
                            print(f"[Milvus] ID={point_id}的向量余弦相似度为{similarity}，低于阈值0.9999")
                            return False
                    else:
                        # 欧几里得距离比较
                        distance = np.sqrt(np.sum((vector - db_values) ** 2))
                        if distance > 0.0001:  # 设置一个接近0的阈值
                            print(f"[Milvus] ID={point_id}的向量欧几里得距离为{distance}，高于阈值0.0001")
                            return False
                except Exception as e:
                    print(f"[Milvus] 比较向量时出错: {e}")
                    return False
            
            print(f"[Milvus] 数据集校验通过，可以重用现有数据")
            return True
            
        except Exception as e:
            print(f"[Milvus] 检查数据集时出错: {e}")
            return False

    def has_same_index(self) -> bool:
        """检查是否有相同参数的索引"""
        try:
            # 检查集合是否存在
            if not utility.has_collection(self.collection_name):
                print(f"集合 {self.collection_name} 不存在")
                return False
                
            # 加载集合
            if not self.collection:
                self.collection = Collection(
                    self.collection_name,
                    consistence_level="STRONG"
                )
                
            # 检查索引是否存在
            try:
                # 明确查询vector字段上的索引
                index_info = self.collection.index(field_name="vector")
                if not index_info:
                    print(f"集合 {self.collection_name} 的vector字段上不存在索引")
                    return False
            except Exception as e:
                print(f"获取索引信息出错，可能是vector字段上不存在索引: {e}")
                return False
                
            # 检查索引类型和参数是否匹配
            index_params = self.get_index_param()
            current_index = index_info.to_dict()
                    
            # 检查索引类型 - 处理不同结构的索引信息
            current_type = None
            if "index_type" in current_index:
                current_type = current_index.get("index_type")
            elif "index_param" in current_index and "index_type" in current_index.get("index_param", {}):
                current_type = current_index.get("index_param", {}).get("index_type")
                
            expected_type = index_params.get("index_type")
            
            if current_type != expected_type:
                print(f"索引类型不匹配: 当前={current_type}, 预期={expected_type}")
                return False
                
            # 检查度量类型 - 处理不同结构的索引信息
            current_metric = None
            if "metric_type" in current_index:
                current_metric = current_index.get("metric_type")
            elif "index_param" in current_index and "metric_type" in current_index.get("index_param", {}):
                current_metric = current_index.get("index_param", {}).get("metric_type")
                
            expected_metric = index_params.get("metric_type")
            
            if current_metric != expected_metric:
                print(f"度量类型不匹配: 当前={current_metric}, 预期={expected_metric}")
                return False
                
            # 检查其他参数 - 处理不同结构的索引信息
            current_params = {}
            if "params" in current_index:
                current_params = current_index.get("params", {})
            elif "index_param" in current_index and "params" in current_index.get("index_param", {}):
                current_params = current_index.get("index_param", {}).get("params", {})
                
            expected_params = index_params.get("params", {})
            
            for key, value in expected_params.items():
                if key not in current_params:
                    print(f"索引参数'{key}'不存在: 预期={value}")
                    return False
                
                if current_params[key] != value:
                    print(f"索引参数'{key}'不匹配: 当前={current_params[key]}, 预期={value}")
                    return False
                    
            print(f"索引校验通过，可以重用现有索引")
            return True

        except Exception as e:
            print(f"检查索引时出错: {e}")
            return False

    def drop(self):
        """Drop existing collection"""
        if utility.has_collection(self.collection_name):
            print(f"[Milvus] collection {self.collection_name} already exists, drop it...")
            utility.drop_collection(self.collection_name)

    def create_collection(self):
        filed_id = FieldSchema(
            name="id",
            dtype=DataType.INT64,
            is_primary=True
        )
        filed_vec = FieldSchema(
            name="vector",
            dtype=DataType.FLOAT_VECTOR,
            dim=self._dim
        )
        schema = CollectionSchema(
            fields=[filed_id, filed_vec],
            description="Test milvus search",
        )
        self.collection = Collection(
            self.collection_name,
            schema,
            consistence_level="STRONG"
        )
        print(f"[Milvus] Create collection {self.collection.describe()} successfully!!!")

    def copy(self, dataset):
        """Copy data to table using batch processing
        Args:
            dataset: h5py Dataset object
        """
        # 创建集合
        self.create_collection()
        
        # 获取训练数据集
        train_data = dataset["train"]
        print(f"[Milvus] Insert {train_data.shape[0]} data into collection {self.collection_name}...")
        
        # 批量处理
        batch_size = 5000
        total_rows = train_data.shape[0]
        print_interval = max(1, total_rows // 10)  # 每10%打印一次
        
        print(f"[Milvus] Start copying {total_rows} rows with batch size {batch_size}...")
        inserted_count = 0
        
        for start_idx in range(0, total_rows, batch_size):
            try:
                end_idx = min(start_idx + batch_size, total_rows)
                batch = train_data[start_idx:end_idx]  # 只读取一部分数据
                batch_size_actual = end_idx - start_idx
                
                entities = [
                    [i for i in range(start_idx, end_idx)],  # id列表
                    batch.tolist()  # 向量数据
                ]
                
                # 插入数据并检查插入结果
                res = self.collection.insert(entities)
                
                # 累计成功插入的行数
                if hasattr(res, "insert_count"):
                    inserted_count += res.insert_count
                else:
                    inserted_count += batch_size_actual
                
                # 每完成10%或最后一批时打印进度
                if end_idx % print_interval == 0 or end_idx == total_rows:
                    print(f"[Milvus] Processed {end_idx}/{total_rows} rows ({(end_idx/total_rows*100):.2f}%)")
            except Exception as e:
                print(f"[Milvus] 批次 {start_idx}-{end_idx} 插入出错: {e}")
                # 继续处理下一批
        
        # 刷新数据确保写入
        self.collection.flush()
        
        # 验证插入结果
        actual_count = self.collection.num_entities
        print(f"[Milvus] 实际插入行数: {actual_count}/{total_rows} ({(actual_count/total_rows*100):.2f}%)")
        
        if actual_count != total_rows:
            print(f"[Milvus] 警告: 插入的行数 ({actual_count}) 与预期不符 ({total_rows})!")
            
        return actual_count

    def get_index_param(self):
        raise NotImplementedError()

    def create_index(self):
        # create index
        if not self.collection:
            self.collection = Collection(
                self.collection_name,
                consistence_level="STRONG"
            )
        
            
        # 需要创建新索引，先释放集合再删除已有索引
        try:
            # 首先释放集合
            print(f"[Milvus] 释放集合 {self.collection_name} 以便修改索引...")
            self.release_collection()
            
            # 获取现有的所有索引
            try:
                indexes = self.collection.index(field_name="vector")
                if indexes:
                    print(f"[Milvus] 删除集合 {self.collection_name} 的vector字段上的现有索引...")
                    # 使用index_name参数替代field_name参数
                    self.collection.drop_index(index_name="vector_index")
                    print(f"[Milvus] 现有索引已删除")
            except Exception as e:
                # 如果获取索引失败，可能是字段上没有索引，或者索引相关的错误
                print(f"[Milvus] 检查vector字段索引时出错: {e}")
                # 仍然尝试删除索引，以防万一
                try:
                    self.collection.drop_index(index_name="vector_index")
                except Exception as e:
                    print(f"[Milvus] 尝试删除索引时出错: {e}")
        except Exception as e:
            print(f"[Milvus] 删除索引时出错: {e}")
        
        print(f"[Milvus] Create index for collection {self.collection_name}...")
        try:
            self.collection.create_index(
                field_name="vector",
                index_params=self.get_index_param(),
                index_name="vector_index",
                timeout=3600*24*7
            )
            utility.wait_for_index_building_complete(
                collection_name=self.collection_name,
                index_name="vector_index",
                timeout=3600*24*7
            )
            index = self.collection.index(field_name="vector")
            index_progress = utility.index_building_progress(
                collection_name=self.collection_name,
                index_name="vector_index",
                timeout=3600*24*7
            )
            print(
                f"[Milvus] Create index {index.to_dict()} {index_progress} for collection {self.collection_name} successfully!!!")
            self.load_collection()
        except Exception as e:
            print(f"[Milvus] 创建索引失败: {e}")
            raise e

    def load_collection(self):
        # load collection
        print(f"[Milvus] Load collection {self.collection_name}...")
        try:
            # 检查集合是否存在
            if not utility.has_collection(self.collection_name):
                print(f"[Milvus] 集合 {self.collection_name} 不存在，无法加载")
                return
                
            # 初始化collection对象（如果为None）
            if not self.collection:
                self.collection = Collection(
                    self.collection_name,
                    consistence_level="STRONG"
                )
                
            # 加载集合
            self.collection.load()
            utility.wait_for_loading_complete(self.collection_name)
            print(f"[Milvus] Load collection {self.collection_name} successfully!!!")
        except Exception as e:
            print(f"[Milvus] 加载集合时出错: {e}")
            # 不抛出异常，以便在多进程环境中继续运行

    def release_collection(self):
        # release collection
        if not self.collection:
            return
            
        print(f"[Milvus] Release collection {self.collection_name}...")
        try:
            self.collection.release()
            print(f"[Milvus] Release collection {self.collection_name} successfully!!!")
        except Exception as e:
            print(f"[Milvus] 释放集合时出错: {e}")
            # 不抛出异常，因为释放失败不应该中断流程

    def query(self, v, n):
        # 确保collection已初始化
        if not self.collection:
            try:
                if utility.has_collection(self.collection_name):
                    self.collection = Collection(
                        self.collection_name,
                        consistence_level="STRONG"
                    )
                    print(f"[Milvus] 初始化集合 {self.collection_name} 成功")
                else:
                    print(f"[Milvus] 集合 {self.collection_name} 不存在，无法查询")
                    return []
            except Exception as e:
                print(f"[Milvus] 初始化集合时出错: {e}")
                return []
                
        try:
            results = self.collection.search(
                data=[v],
                anns_field="vector",
                param=self.search_params,
                limit=n,
                output_fields=["id"]
            )
            ids = [r.entity.get("id") for r in results[0]]
            return ids
        except Exception as e:
            print(f"[Milvus] 查询时出错: {e}")
            return []
    
    def get_memory_usage(self):
        return 0
    
    def get_table_usage(self):
        return 0

class MilvusFLAT(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self.name = f"MilvusFLAT metric:{self._metric}"

    def get_index_param(self):
        return {
            "index_type": "FLAT",
            "metric_type": self._metric_type,
            "params": {}
        }

    def set_query_arguments(self):
        self.search_params = {
            "metric_type": self._metric_type
        }
        self.name = f"milvus-flat metric:{self._metric}"


class MilvusIVFSQ8(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_nlist = index_param.get("nlist", None)
        self.name = f"MilvusIVFSQ8 metric:{self._metric}"

    def get_index_param(self):
        return {
            "index_type": "IVF_SQ8",
            "metric_type": self._metric_type,
            "params": {
                "nlist": self._index_nlist
            }
        }

    def set_query_arguments(self, nprobe):
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"nprobe": nprobe}
        }
        self.name = f"milvus-ivfsq8 metric:{self._metric}, index_nlist:{self._index_nlist}, search_nprobe:{nprobe}"


class MilvusIVFFLAT(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_nlist = index_param.get("nlist", None)
        self.name = f"MilvusIVFFLAT metric:{self._metric}"

    def get_index_param(self):
        return {
            "index_type": "IVF_FLAT",
            "metric_type": self._metric_type,
            "params": {
                "nlist": self._index_nlist
            }
        }

    def set_query_arguments(self, nprobe):
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"nprobe": nprobe}
        }
        self.name = f"milvus-ivfflat metric:{self._metric}, index_nlist:{self._index_nlist}, search_nprobe:{nprobe}"
    


class MilvusIVFPQ(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_nlist = index_param.get("nlist", None)
        self._index_m = index_param.get("m", None)
        self._index_nbits = index_param.get("nbits", None)
        self.name = f"MilvusIVFPQ metric:{self._metric}"

    def get_index_param(self):
        assert self._dim % self._index_m == 0, "dimension must be able to be divided by m"
        return {
            "index_type": "IVF_PQ",
            "metric_type": self._metric_type,
            "params": {
                "nlist": self._index_nlist,
                "m": self._index_m,
                "nbits": self._index_nbits if self._index_nbits else 8
            }
        }

    def set_query_arguments(self, nprobe):
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"nprobe": nprobe}
        }
        self.name = f"milvus-ivfpq metric:{self._metric}, index_nlist:{self._index_nlist}, search_nprobe:{nprobe}"


class MilvusHNSW(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_m = index_param.get("M", None)
        self._index_ef = index_param.get("efConstruction", None)
        self.name = f"MilvusHNSW metric:{self._metric}"

    def get_index_param(self):
        return {
            "index_type": "HNSW",
            "metric_type": self._metric_type,
            "params": {
                "M": self._index_m,
                "efConstruction": self._index_ef
            }
        }

    def set_query_arguments(self, ef):
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"ef": ef}
        }
        self.name = f"milvus-hnsw metric:{self._metric}, index_M:{self._index_m}, index_ef:{self._index_ef}, search_ef={ef}"


class MilvusSCANN(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_nlist = index_param.get("nlist", None)
        self.name = f"MilvusSCANN metric:{self._metric}"

    def get_index_param(self):
        return {
            "index_type": "SCANN",
            "metric_type": self._metric_type,
            "params": {
                "nlist": self._index_nlist
            }
        }

    def set_query_arguments(self, nprobe):
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {"nprobe": nprobe}
        }
        self.name = f"milvus-scann metric:{self._metric}, index_nlist:{self._index_nlist}, search_nprobe:{nprobe}"
